import cv2
import numpy as np
import matplotlib.pyplot as plt
import cmath

def SeparateOddEven(array):
	if len(array) > 4:
		return SeparateOddEven(array[::2]) + SeparateOddEven(array[1::2])
	elif len(array) == 4:
		return array[::2] + array[1::2]


def RearrangeIndex(n):
	return SeparateOddEven(range(n))


# size of lookup table can be cut half
def GenerateTwiddleFactorLookup(N):
	lookup = [1j]*N
	middle = N/2 
	for i in range(middle):
		lookup[i] = cmath.e**(-1j*2*cmath.pi*i/N)
		lookup[i+middle] = -1 * lookup[i]

	return lookup


def SubgroupFFT(array, size, N, TFLookup):
	intermedia = np.array([1]*N, dtype=complex)
	middle = size / 2

	for base in range(0, N, size):
		for i in range(middle):
			a = array[base + i] * 1.0
			b = array[base + i + middle] * TFLookup[i*N/size]  #* cmath.e**(-1j * 2 * cmath.pi * 0/N)
			intermedia[base + i] = a + b
			intermedia[base + i + middle] = a - b

	return intermedia


def FFT1D(array, N, lookup):
	"DIT-FFT"
	indexes = RearrangeIndex(N)

	intermedia = np.array([1]*N, dtype=complex)

	for i in range(0, N, 2):
		a = array[indexes[i]] * 1.0
		b = array[indexes[i+1]] * 1.0 # * cmath.e**(-1j * 2 * cmath.pi * 0/N)
		intermedia[i] = a + b
		intermedia[i+1] = a - b


	subgroupLength = 4
	while subgroupLength <= N:
		intermedia = SubgroupFFT(intermedia, subgroupLength, N, lookup)
		subgroupLength *= 2


	return intermedia


def FFT2D(matrix, rows, cols):
	result = np.empty([rows, cols], dtype=complex)
	lookup = GenerateTwiddleFactorLookup(cols)

	for i in range(rows):
		result[i, :] = FFT1D(matrix[i, :], cols, lookup)

	for j in range(cols):
		result[:, j] = FFT1D(result[:, j], rows, lookup)


	return result


#
# IFFT
#

def IFFT2D(matrix, rows, cols):
	tmp = np.conjugate(matrix)

	result = np.empty([rows, cols], dtype=complex)
	lookup = GenerateTwiddleFactorLookup(cols)

	for i in range(rows):
		result[i, :] = np.conjugate(FFT1D(matrix[i, :], cols, lookup)) / cols

	for j in range(cols):
		result[:, j] = np.conjugate(FFT1D(result[:, j], rows, lookup)) / rows

	# ignore the image component which is caused by percision issue
	# meanwhile flip the image horizontal
	realResult = np.empty([rows, cols], dtype=np.uint8)
	for i in range(rows):
		realResult[i, :]  = result[i, :][::-1]

	return realResult


if __name__ == '__main__':
	print SeparateOddEven(range(8))
	print RearrangeIndex(8)
	print SeparateOddEven(range(16))
	#print SeparateOddEven(range(256))
	print GenerateTwiddleFactorLookup(8)
	print "=========FFT Result Compare==============="
	N = 64
	print "From Numpy"
	print np.fft.fft(np.sin(np.arange(N)))
	print "From My Code"
	lookup = GenerateTwiddleFactorLookup(N)
	fft = FFT1D(np.sin(np.arange(N)), N, lookup)
	print fft
	print "Raw Data"
	print np.sin(np.arange(N))
	print "Numpy Inverse Transform by FFT"
	print np.fft.fft(np.conjugate(fft))
	print "Numpy Inverse Transform by API"
	print np.fft.ifft(fft)
	print "Inverse Transform"
	myInput = FFT1D(np.conjugate(fft), N, lookup)
	print myInput
	print "=========IFFT Result Compare==============="
	N = 4
	print np.fft.ifft([-2j, 4, 1j, 0])
	print np.conjugate(np.fft.fft(np.conjugate([-2j, 4, 1j, 0]))) / N
	lookup = GenerateTwiddleFactorLookup(N)
	print np.conjugate(FFT1D(np.conjugate(np.array([-2j, 4, 1j, 0])), N, lookup)) / N
	print "=========FFT/IFFT on image==============="

	rawImg = cv2.imread('./img/fig5.jpg',0)
	
	height, width = rawImg.shape

	f = np.fft.fft2(rawImg)
	fshift = np.fft.fftshift(f)
	magnitude_spectrum = 20*np.log(np.abs(fshift))

	plt.subplot(121),plt.imshow(magnitude_spectrum, cmap = 'gray')
	plt.title('Numpy Magnitude Spectrum'), plt.xticks([]), plt.yticks([])

	print "Calculating DIT FFT 2D"
	myf = FFT2D(rawImg, height, width)
	print "Finish calculating DIT FFT 2D"
	
	myfshift = np.fft.fftshift(myf)
	my_magnitude_spectrum = 20*np.log(np.abs(myfshift))
	plt.subplot(122),plt.imshow(my_magnitude_spectrum, cmap = 'gray')
	plt.title('My Magnitude Spectrum'), plt.xticks([]), plt.yticks([])

	print "Calculating DIT IFFT 2D"
	myImg = IFFT2D(myf, height, width)
	print "Finish calculating DIT IFFT 2D"

	cv2.namedWindow("Original Image", cv2.CV_WINDOW_AUTOSIZE)
	cv2.imshow("Original Image", rawImg)
	cv2.namedWindow("Image from IFFT", cv2.CV_WINDOW_AUTOSIZE)
	cv2.imshow("Image from IFFT", myImg)

	plt.show()
	
	cv2.waitKey(0)
	cv2.destroyAllWindows()